import torch
from PIL import Image
import torchvision.transforms as transforms
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32*16*16, 64),
            nn.ReLU(),
            nn.Linear(64, 2)  # 二分类
        )
    
    def forward(self, x):
        x = self.conv(x)
        x = self.fc(x)
        return x

# 1. 加载训练好的模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 假设你已经训练了一个模型，并保存在 'data/models/simple_cnn.pth' 路径下
model = SimpleCNN().to(device)
model.load_state_dict(torch.load('E:/My-PyTorch-Learn/data/models/simple_cnn.pth'))  # 修改为绝对路径 # 加载模型权重
model.eval()  # 切换到评估模式

# 2. 读取单张测试图片
img_path = 'E:/My-PyTorch-Learn/data/train_data/{}/0.png'
# 标准输出流读取输入，组合图片路径
img_path = img_path.format(input())
image = Image.open(img_path)

# 3. 数据预处理：调整图片大小，转换为 tensor
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 调整为 64x64 尺寸
    transforms.ToTensor(),        # 转换为 Tensor 类型
])

image = transform(image).unsqueeze(0)  # 增加 batch 维度

# 4. 使用模型进行推理
image = image.to(device)  # 将图像数据移动到正确的设备（GPU/CPU）
with torch.no_grad():  # 不需要计算梯度，推理时加速
    output = model(image)
    pred = torch.argmax(output, 1)  # 获取最大输出的类别（0 或 1）

# 5. 打印预测结果
result = pred.item();# 1=wrong, 0=correct
if(result==1):
    print("预测标签: wrong 含有杂质")
elif(result==0):
    print("预测标签: correct 不含杂质")
